#ifndef __ASCEND_C_TENSOR__H__
#define __ASCEND_C_TENSOR__H__
#include "common.h"
#include <cstdint>
namespace AscendC
{
enum OpType
{
    OP_None = 0,
    // 矢量单目
    OP_Exp = 101,
    OP_Ln = 102,
    OP_Abs = 103,
    OP_Reciprocal = 104, // y = 1 / x
    OP_Sqrt = 105,
    OP_Rsqrt = 106, // y = 1 / sqrt(x)
    OP_Not = 107,
    OP_Relu = 108,
    // 矢量双目
    OP_Add = 201,
    OP_Sub = 202,
    OP_Mul = 203,
    OP_Div = 204,
    OP_Max = 205,
    OP_Min = 206,
    OP_And = 207,
    OP_Or = 208,
    OP_AddRelu = 209,
    // 标量双目
    OP_Adds = 301,
    OP_Muls = 302,
    OP_Maxs = 303,
    OP_Mins = 304,
    OP_ShiftLeft = 305,
    OP_ShiftRight = 306,
    OP_LeakyRelu = 307,
    // 标量三目
    OP_Axpy = 401, // y = x * scalar + y
    // 比较指令
    OP_Compare = 501,
    // 选择指令
    OP_Select = 601,
    OP_GatherMask = 602,
    // 精度转换指令
    OP_Cast = 701,
    // 归约指令
    OP_ReduceMax = 801,
    OP_ReduceMin = 802,
    OP_ReduceSum = 803,
    OP_WholeReduce = 804,
    OP_BlockReduceMax = 805,
    OP_BlockReduceMin = 806,
    OP_BlockReduceSum = 807,
    OP_PairReduce = 808, // 相邻两个（奇偶）元素求和，例如（a1, a2, a3, a4, a5, a6...），相邻两个数据求和为（a1+a2,
                         // a3+a4, a5+a6, ......）。
    // 矩阵计算
    OP_LoadData = 901,
    OP_Mmad = 902,
    OP_Gemm = 903,
    // 数据搬运
    OP_DataCopy = 1001,
    //
};

union IntArg {
    SELMODE selfMode;
    CMPMODE cmpMode;
    RoundMode roundMode;
};

template <typename DST, typename ARG, typename EXTRA> class OpArgs
{
  public:
    OpType op;
    LocalTensor<DST> dst;
    LocalTensor<ARG> arg1;
    LocalTensor<ARG> arg2;
    LocalTensor<EXTRA> arg3;
    ARG scalar_arg1;
    IntArg mode;
    uint32_t calcCount;
    bool boolArg = false;
};

template <typename DST, typename ARG, typename EXTRA> void run_op(OpArgs<DST, ARG, EXTRA> args)
{
    switch (args.op)
    {
    // 矢量单目运算符
    case OP_Exp: {
        Exp(args.dst, args.arg1, args.calcCount);
    }
    break;
    case OP_Ln: {
        Ln(args.dst, args.arg1, args.calcCount);
    }
    break;
    case OP_Abs: {
        Abs(args.dst, args.arg1, args.calcCount);
    }
    break;
    case OP_Reciprocal: {
        Reciprocal(args.dst, args.arg1, args.calcCount);
    }
    break;
    case OP_Sqrt: {
        Sqrt(args.dst, args.arg1, args.calcCount);
    }
    break;
    case OP_Rsqrt: {
        Rsqrt(args.dst, args.arg1, args.calcCount);
    }
    break;
    case OP_Not: {
        Not(args.dst, args.arg1, args.calcCount);
    }
    break;
    case OP_Relu: {
        Relu(args.dst, args.arg1, args.calcCount);
    }
    break;
    // 矢量双目运算符
    case OP_Add: {
        Add(args.dst, args.arg1, args.arg2, args.calcCount);
    }
    break;
    case OP_Sub: {
        Sub(args.dst, args.arg1, args.arg2, args.calcCount);
    }
    break;
    case OP_Mul: {
        Sub(args.dst, args.arg1, args.arg2, args.calcCount);
    }
    break;
    case OP_Div: {
        Div(args.dst, args.arg1, args.arg2, args.calcCount);
    }
    break;
    case OP_Max: {
        Max(args.dst, args.arg1, args.arg2, args.calcCount);
    }
    break;
    case OP_Min: {
        Min(args.dst, args.arg1, args.arg2, args.calcCount);
    }
    break;
    case OP_And: {
        And(args.dst, args.arg1, args.arg2, args.calcCount);
    }
    break;
    case OP_Or: {
        Or(args.dst, args.arg1, args.arg2, args.calcCount);
    }
    break;
    // 标量双目
    case OP_Adds: {
        Adds(args.dst, args.arg1, args.scalar_arg1, args.calcCount);
    }
    break;
    case OP_Muls: {
        Muls(args.dst, args.arg1, args.scalar_arg1, args.calcCount);
    }
    break;
    case OP_Maxs: {
        Maxs(args.dst, args.arg1, args.scalar_arg1, args.calcCount);
    }
    break;
    case OP_Mins: {
        Mins(args.dst, args.arg1, args.scalar_arg1, args.calcCount);
    }
    break;
    case OP_ShiftLeft: {
        ShiftLeft(args.dst, args.arg1, args.scalar_arg1, args.calcCount);
    }
    break;
    case OP_ShiftRight: {
        ShiftRight(args.dst, args.arg1, args.scalar_arg1, args.calcCount);
    }
    break;
    case OP_LeakyRelu: {
        LeakyRelu(args.dst, args.arg1, args.scalar_arg1, args.calcCount);
    }
    break;
    // 标量三目
    case OP_Axpy: {
        Axpy(args.dst, args.arg1, args.scalar_arg1, args.calcCount);
    }
    break;
    // 比较
    case OP_Compare: {
        Compare(args.dst, args.arg1, args.arg2, args.mode.cmpMode, args.calcCount);
    }
    break;
    // 选择指令
    case OP_Select: {
        switch (args.mode.cmpMode)
        {
        case SELMODE::VSEL_CMPMASK_SPR:
        case SELMODE::VSEL_TENSOR_TENSOR_MODE: {
            Select(args.dst, args.arg3, args.arg1, args.arg2, args.mode.cmpMode, args.calcCount);
        }
        break;
        case SELMODE::VSEL_TENSOR_SCALAR_MODE: {
            Select(args.dst, args.arg3, args.arg1, args.scalar_arg1, SELMODE::VSEL_TENSOR_SCALAR_MODE, args.calcCount);
        }
        break;
        }
    }
    break;
    // 归约指令
    case OP_ReduceMax: {
        ReduceMax(args.dst, args.arg1, args.arg2, args.calcCount, args.boolArg);
    }
    break;
    case OP_ReduceMin: {
        ReduceMin(args.dst, args.arg1, args.arg2, args.calcCount, args.boolArg);
    }
    break;
    case OP_ReduceSum: {
        ReduceSum(args.dst, args.arg1, args.arg2, args.calcCount, args.boolArg);
    }
    break;
    // 数据搬运
    case OP_DataCopy: {
        DataCopy(args.dst, args.arg1, args.calcCount);
    }
    break;
    }
}
#define DEF_DOUBLE_OP_FUNC(op, fn_name)                                                                                \
    __aicore__ inline Tensor fn_name(Tensor &other)                                                                    \
    {                                                                                                                  \
        auto c = MIN(calcCount, other.calcCount);                                                                      \
        auto temp = getTemp();                                                                                         \
        op(temp, data, other.data, c);                                                                                 \
        return Tensor(temp, c);                                                                                        \
    }                                                                                                                  \
    __aicore__ inline Tensor fn_name(LocalTensor<T> &other)                                                            \
    {                                                                                                                  \
        auto temp = getTemp();                                                                                         \
        op(temp, data, other, calcCount);                                                                              \
        return Tensor(temp, calcCount);                                                                                \
    }                                                                                                                  \
    __aicore__ inline Tensor &fn_name##_(Tensor &other)                                                                \
    {                                                                                                                  \
        auto c = MIN(calcCount, other.calcCount);                                                                      \
        op(data, data, other.data, c);                                                                                 \
        return *this;                                                                                                  \
    }                                                                                                                  \
    __aicore__ inline Tensor &fn_name##_(LocalTensor<T> &other)                                                        \
    {                                                                                                                  \
        op(data, data, other, calcCount);                                                                              \
        return *this;                                                                                                  \
    }

#define DEF_DOUBLE_OP_SCALAR_FUNC(op, fn_name)                                                                         \
    __aicore__ inline Tensor fn_name(T other)                                                                          \
    {                                                                                                                  \
        auto temp = getTemp();                                                                                         \
        op(temp, data, other, calcCount);                                                                              \
        return Tensor(temp, calcCount);                                                                                \
    }                                                                                                                  \
    __aicore__ inline Tensor &fn_name##_(T other)                                                                      \
    {                                                                                                                  \
        op(data, data, other, calcCount);                                                                              \
        return *this;                                                                                                  \
    }

#define DEF_DOUBLE_OP(op, cop)                                                                                         \
    __aicore__ inline Tensor operator cop(Tensor &other)                                                               \
    {                                                                                                                  \
        auto c = MIN(calcCount, other.calcCount);                                                                      \
        auto temp = getTemp();                                                                                         \
        op(temp, data, other.data, c);                                                                                 \
        return Tensor(temp, c);                                                                                        \
    }                                                                                                                  \
    __aicore__ inline Tensor operator cop(LocalTensor<T> &other)                                                       \
    {                                                                                                                  \
        auto temp = getTemp();                                                                                         \
        op(temp, data, other, calcCount);                                                                              \
        return Tensor(temp, calcCount);                                                                                \
    }                                                                                                                  \
    __aicore__ inline Tensor &operator cop##=(Tensor &other)                                                           \
    {                                                                                                                  \
        auto c = MIN(calcCount, other.calcCount);                                                                      \
        op(data, data, other.data, c);                                                                                 \
        return *this;                                                                                                  \
    }                                                                                                                  \
    __aicore__ inline Tensor &operator cop##=(LocalTensor<T> &other)                                                   \
    {                                                                                                                  \
        op(data, data, other, calcCount);                                                                              \
        return *this;                                                                                                  \
    }

#define DEF_DOUBLE_OP_SCALAR(op, cop)                                                                                  \
    __aicore__ inline Tensor operator cop(T other)                                                                     \
    {                                                                                                                  \
        auto temp = getTemp();                                                                                         \
        op(temp, data, other, calcCount);                                                                              \
        return Tensor(temp, calcCount);                                                                                \
    }                                                                                                                  \
    __aicore__ inline Tensor operator cop##=(T other)                                                                  \
    {                                                                                                                  \
        op(data, data, other, calcCount);                                                                              \
        return *this;                                                                                                  \
    }

#define DEF_DOUBLE_OP_WITH_SCALAR(op, cop)                                                                             \
    DEF_DOUBLE_OP(op, cop)                                                                                             \
    DEF_DOUBLE_OP_SCALAR(op##s, cop)

#define DEF_DOUBLE_OP_FUNC_WITH_SCALAR(op, fn_name)                                                                    \
    DEF_DOUBLE_OP_FUNC(op, fn_name)                                                                                    \
    DEF_DOUBLE_OP_SCALAR_FUNC(op##s, fn_name)

#define DEF_COMPARE_OPERATOR(mode, cop)                                                                                \
    __aicore__ inline Tensor<uint8_t> operator cop(LocalTensor<T> &other)                                              \
    {                                                                                                                  \
        auto temp = getTemp();                                                                                         \
        auto int8_temp = temp.template ReinterpretCast<uint8_t>();                                                     \
        Compare(int8_temp, data, other, mode, calcCount);                                                              \
        return Tensor<uint8_t>(int8_temp, calcCount);                                                                  \
    }                                                                                                                  \
    __aicore__ inline Tensor<uint8_t> operator cop(Tensor &other)                                                      \
    {                                                                                                                  \
        auto temp = getTemp();                                                                                         \
        auto int8_temp = temp.template ReinterpretCast<uint8_t>();                                                     \
        auto c = MIN(calcCount, other.calcCount);                                                                      \
        Compare(int8_temp, data, other.data, mode, c);                                                                 \
        return Tensor<uint8_t>(int8_temp, c);                                                                          \
    }

#define DEF_SINGLE_OPERATOR(op, fn_name)                                                                               \
    __aicore__ inline Tensor &fn_name()                                                                                \
    {                                                                                                                  \
        op(data, data, calcCount);                                                                                     \
        return *this;                                                                                                  \
    }

#define DEF_MY_TENSOR_TEMP                                                                                             \
    DefBufVECCALC(TENSOR_TEMP_1);                                                                                      \
    DefBufVECCALC(TENSOR_TEMP_2);                                                                                      \
    uint32_t tensor_temp_index;
#define TBUF_GET_MY_TENSOR_TEMP                                                                                        \
    TBufGet(TENSOR_TEMP_1, TypeOf(X));                                                                                 \
    TBufGet(TENSOR_TEMP_2, TypeOf(X));
#define INIT_MY_TENSOR_TEMP_AND_GET(length, dtype)                                                                     \
    InitTBufBuffer(TENSOR_TEMP_1, length * sizeof(dtype));                                                             \
    InitTBufBuffer(TENSOR_TEMP_2, length * sizeof(dtype));                                                             \
    TBUF_GET_MY_TENSOR_TEMP
#define INIT_MY_TENSOR_TEMP(x)                                                                                         \
    x.temp1 = TBufTensorName(TENSOR_TEMP_1);                                                                           \
    x.temp2 = TBufTensorName(TENSOR_TEMP_2);                                                                           \
    x.temp_index = &tensor_temp_index;
#define DEFINE_MY_TENSOR(name, dtype, lt, c)                                                                     \
    Tensor<dtype> name;                                                                                                \
    name.data = lt;                                                                                           \
    name.calcCount = c;                                                                                        \
    name.temp1 = TBufTensorName(TENSOR_TEMP_1).template ReinterpretCast<dtype>();                                      \
    name.temp2 = TBufTensorName(TENSOR_TEMP_2).template ReinterpretCast<dtype>();                                      \
    name.temp_index = &tensor_temp_index;
#define MY_TENSOR_UPDATE(name, d, c)                                                                                      \
    name.data = d;                                                                                           \
    name.calcCount = c;

template <typename T> class Tensor
{
  public:
    LocalTensor<T> temp1;
    LocalTensor<T> temp2;
    uint32_t *temp_index = 0;
    LocalTensor<T> data;
    uint32_t calcCount = 0;

  public:
    __aicore__ inline Tensor()
    {
    }

    template <typename D> __aicore__ inline Tensor(D v)
    {
        this->data = v;
        this->calcCount = this->data.GetSize();
    }

    template <typename D> __aicore__ inline Tensor(D v, D temp1, D temp2, uint32_t *temp_index)
    {
        this->data = v;
        this->temp1 = temp1;
        this->temp2 = temp2;
        this->temp_index = temp_index;
        this->calcCount = this->data.GetSize();
    }

    template <typename D> __aicore__ inline Tensor(D v, uint32_t calcCount)
    {
        this->data = v;
        this->calcCount = calcCount;
    }

    template <typename D> __aicore__ inline Tensor(D v, uint32_t calcCount, D temp1, D temp2, uint32_t *temp_index)
    {
        this->data = v;
        this->temp1 = temp1;
        this->temp2 = temp2;
        this->temp_index = temp_index;
        this->calcCount = calcCount;
    }

    template <typename D> __aicore__ inline Tensor<D> getRet()
    {
        if (*temp_index == 0)
        {
            return Tensor<D>(temp1.template ReinterpretCast<D>(), calcCount);
        }
        return Tensor<D>(temp2.template ReinterpretCast<D>(), calcCount);
    }

    __aicore__ inline auto getTemp(bool exchange = true)
    {
        auto ret = *temp_index % 2 == 0 ? temp1 : temp2;
        if (exchange)
        {
            *temp_index = (*temp_index + 1) % 2;
        }
        return ret;
    }

    template <typename D> __aicore__ inline Tensor<D> reinterpret()
    {
        return Tensor<D>(data.template ReinterpretCast<D>(), calcCount);
    }

    // 单目运算符
    __aicore__ inline Tensor &operator=(Tensor other)
    {
        DataCopy(this->data, other.data, other.calcCount);
        this->calcCount = other.calcCount;
        return *this;
    }

    __aicore__ inline Tensor &operator=(LocalTensor<T> &v)
    {
        DataCopy(this->data, v, calcCount);
        return *this;
    }

    __aicore__ inline Tensor &operator=(T v)
    {
        for (auto i = 0; i < data.GetSize(); ++i)
        {
            data.SetValue(i, v);
        }
        return *this;
    }
    DEF_SINGLE_OPERATOR(Exp, exp);
    DEF_SINGLE_OPERATOR(Ln, ln);
    DEF_SINGLE_OPERATOR(Abs, abs);
    DEF_SINGLE_OPERATOR(Reciprocal, reciprocal);
    DEF_SINGLE_OPERATOR(Sqrt, sqrt);
    DEF_SINGLE_OPERATOR(Rsqrt, rsqrt);
    DEF_SINGLE_OPERATOR(Not, operator~);
    DEF_SINGLE_OPERATOR(Relu, relu);

    // 双目运算符
    DEF_DOUBLE_OP_WITH_SCALAR(Add, +);
    DEF_DOUBLE_OP_WITH_SCALAR(Mul, *);
    DEF_DOUBLE_OP(Sub, -);
    DEF_DOUBLE_OP(Div, /);
    DEF_DOUBLE_OP(And, &);
    DEF_DOUBLE_OP(Or, |);
    DEF_DOUBLE_OP_FUNC(AddRelu, addRelu);
    DEF_DOUBLE_OP_FUNC_WITH_SCALAR(Max, max);
    DEF_DOUBLE_OP_FUNC_WITH_SCALAR(Min, min);
    DEF_DOUBLE_OP_SCALAR(ShiftLeft, <<);
    DEF_DOUBLE_OP_SCALAR(ShiftRight, >>);
    DEF_DOUBLE_OP_SCALAR_FUNC(LeakyRelu, leakyRelu);

    __aicore__ inline Tensor operator-(half other)
    {
        auto temp = getTemp();
        Adds(temp, data, T(-float(other)), calcCount);
        return Tensor(temp, calcCount);
    }

    __aicore__ inline Tensor operator-(int16_t other)
    {
        auto temp = getTemp();
        Adds(temp, data, T(-other), calcCount);
        return Tensor(temp, calcCount);
    }

    __aicore__ inline Tensor operator-(int32_t other)
    {
        auto temp = getTemp();
        Adds(temp, data, T(-other), calcCount);
        return Tensor(temp, calcCount);
    }

    __aicore__ inline Tensor operator-(float other)
    {
        auto temp = getTemp();
        Adds(temp, data, T(-other), calcCount);
        return Tensor(temp, calcCount);
    }

    __aicore__ inline Tensor operator/(half other)
    {
        auto temp = getTemp();
        Muls(temp, data, T(1 / float(other)), calcCount);
        return Tensor(temp, calcCount);
    }

    __aicore__ inline Tensor operator/(float other)
    {
        auto temp = getTemp();
        Muls(temp, data, T(1 / other), calcCount);
        return Tensor(temp, calcCount);
    }

    __aicore__ inline Tensor operator-=(half other)
    {
        Adds(data, data, T(-float(other)), calcCount);
        return *this;
    }

    __aicore__ inline Tensor operator-=(int16_t other)
    {
        Adds(data, data, -other, calcCount);
        return *this;
    }

    __aicore__ inline Tensor operator-=(float other)
    {
        Adds(data, data, -other, calcCount);
        return *this;
    }

    __aicore__ inline Tensor operator-=(int32_t other)
    {
        Adds(data, data, -other, calcCount);
        return *this;
    }

    __aicore__ inline Tensor operator/=(T other)
    {
        Muls(data, data, 1 / other, calcCount);
        return *this;
    }

    DEF_COMPARE_OPERATOR(CMPMODE::EQ, ==);
    DEF_COMPARE_OPERATOR(CMPMODE::NE, !=);
    DEF_COMPARE_OPERATOR(CMPMODE::GT, >);
    DEF_COMPARE_OPERATOR(CMPMODE::LT, <);
    DEF_COMPARE_OPERATOR(CMPMODE::GE, >=);
    DEF_COMPARE_OPERATOR(CMPMODE::LE, <=);

    // 三目运算符
    __aicore__ inline Tensor &apxy(LocalTensor<T> &v, T scalar_value)
    {
        Axpy(data, v, scalar_value, calcCount);
        return *this;
    }

    __aicore__ inline Tensor &apxy(Tensor<T> &v, T scalar_value)
    {
        auto c = MIN(calcCount, v.calcCount);
        Axpy(data, v.data, scalar_value, c);
        return *this;
    }

    // 选择指令
    template <typename D> __aicore__ inline auto select(D v, T else_value = 0)
    {
        auto c = MIN(calcCount, v.calcCount);
        auto temp = getTemp();
        Select(temp, v.data, data, else_value, SELMODE::VSEL_TENSOR_SCALAR_MODE, c);
        return Tensor(temp, c);
    }

    template <typename D> __aicore__ inline auto select_(D v, T else_value = 0)
    {
        auto c = MIN(calcCount, v.calcCount);
        Select(data, v.data, data, else_value, SELMODE::VSEL_TENSOR_SCALAR_MODE, c);
        return *this;
    }
};
} // namespace AscendC

#endif